import numpy as np
import tool
from mapping import get_type
import mapping
import tool
import datetime


def cal_product_daily_quantity():
    data = tool.get_np("2x.csv")
    unique_dates_codes = np.unique(data[["date", "code"]], axis=0)
    result = np.zeros(
        len(unique_dates_codes),
        dtype=[("date", "U10"), ("code", "i4"), ("total_quantity", "f8")],
    )
    for i, (date, code) in enumerate(unique_dates_codes):
        mask = (data["date"] == date) & (data["code"] == code)
        total_quantity = np.sum(data["quantity"][mask])
        result[i] = (date, code, total_quantity)

    print(result[:30])
    return result


def cal_type_daily_quantity():
    data = tool.get_np("product_daily_quantity.csv")
    product_types = np.array([get_type(row["code"]) for row in data])
    result_dtype = [("date", "U10"), ("type", "i4"), ("total_sales", "f8")]
    result = []
    unique_dates = np.unique(data["date"])
    for date in unique_dates:
        mask = data["date"] == date
        daily_data = data[mask]
        daily_types = product_types[mask]
        unique_types = np.unique(daily_types)
        for t in unique_types:
            type_mask = daily_types == t
            total_sales = np.sum(daily_data["total_quantity"][type_mask])
            result.append((date, t, total_sales))
            
    result_array = np.array(result, dtype=result_dtype)
    print(result_array[:20])
    return result_array


def split_by_type():
    data = tool.get_np("type_daily_quantity.csv")
    unique_types = np.unique(data["type"])
    for t in unique_types:
        type_mask = data["type"] == t
        type_data = data[type_mask]
        dtype = [
            ("date", "U10"),
            ("total_quantity", "f8"),
        ]
        result_data = np.empty(len(type_data), dtype=dtype)
        for i, row in enumerate(type_data):
            result_data[i]["date"] = row["date"]
            result_data[i]["total_quantity"] = row["total_quantity"]
        output_file = f"type_daily_total/{t}.csv"
        tool.get_csv(result_data, output_file)
        print(f"Saved {output_file}")


def cal_monthly_total(data):
    months = np.array(
        [
            datetime.datetime.strptime(d, "%Y-%m-%d").strftime("%Y-%m")
            for d in data["date"]
        ]
    )
    unique_months = np.unique(months)
    monthly_totals = []
    for month in unique_months:
        mask = months == month
        total_sales = np.sum(data["total_quantity"][mask])
        monthly_totals.append((month, total_sales))
    result_dtype = [("month", "U7"), ("monthly_total", "f8")]
    return np.array(monthly_totals, dtype=result_dtype)


def split_by_month(source_csv, destination_csv):
    data = tool.get_np(source_csv)
    result_data = cal_monthly_total(data)
    tool.get_csv(result_data, destination_csv)


def fill_missing_dates(data):
    dates = np.array(
        [datetime.datetime.strptime(d, "%Y-%m-%d").date() for d in data["date"]]
    )
    start_date = datetime.date(2020, 7, 1)
    end_date = datetime.date(2023, 6, 30)
    all_dates = np.array(
        [
            start_date + datetime.timedelta(days=i)
            for i in range((end_date - start_date).days + 1)
        ]
    )
    date_to_quantity = {
        date: total_quantity
        for date, total_quantity in zip(dates, data["total_quantity"])
    }
    filled_data = []
    for date in all_dates:
        total_quantity = date_to_quantity.get(date, 0.0)  # 如果日期不存在，补0
        filled_data.append((date.strftime("%Y-%m-%d"), total_quantity))
    result_dtype = [("date", "U10"), ("total_quantity", "f8")]
    return np.array(filled_data, dtype=result_dtype)


def split_by_product():
    data = tool.get_np("product_daily_quantity.csv")
    unique_types = np.unique(data["code"])
    for t in unique_types:
        type_mask = data["code"] == t
        type_data = data[type_mask]
        dtype = [
            ("date", "U10"),
            ("total_quantity", "f8"),
        ]
        result_data = np.empty(len(type_data), dtype=dtype)
        for i, row in enumerate(type_data):
            result_data[i]["date"] = row["date"]
            result_data[i]["total_quantity"] = row["total_quantity"]
        result_data1 = fill_missing_dates(result_data)
        output_file = f"product_daily_total/{t}.csv"
        tool.get_csv(result_data1, output_file)
        print(f"Saved {output_file}")


def cal_monthly_total(data):
    months = np.array(
        [
            datetime.datetime.strptime(d, "%Y-%m-%d").strftime("%Y-%m")
            for d in data["date"]
        ]
    )
    unique_months = np.unique(months)
    monthly_totals = []
    for month in unique_months:
        mask = months == month
        total_sales = np.sum(data["total_quantity"][mask])
        monthly_totals.append((month, total_sales))
    result_dtype = [("month", "U7"), ("monthly_total", "f8")]
    return np.array(monthly_totals, dtype=result_dtype)


def split_by_month_product():
    data1 = tool.get_np("product_daily_quantity.csv")
    product = np.unique(data1["code"])
    for i in product:
        data = tool.get_np(f"product_daily_total/{i}.csv")
        result_data = cal_monthly_total(data)
        tool.get_csv(result_data, f"product_monthly_total/{i}.csv")


def cal_daily_total():
    data = tool.get_np("type_daily_quantity.csv")
    unique_dates = np.unique(data["date"])
    result_list = []
    for date in unique_dates:
        tmask = data["date"] == date
        total_quantity = np.sum(data["total_quantity"][tmask])
        result_list.append((date, total_quantity))
    dtype = [("date", "U10"), ("total_quantity", "f8")]
    result_array1 = np.array(result_list, dtype=dtype)
    result_array = fill_missing_dates(result_array1)
    return result_array


def cal_product_total():
    data = tool.get_np("product_daily_quantity.csv")
    result_list = []
    for i in range(1, max(data["code"]) + 1):
        tmask = data["code"] == i
        total_quantity = np.sum(data["total_quantity"][tmask])
        result_list.append((i, total_quantity))
    dtype = [("code", "i4"), ("total_quantity", "f8")]
    return np.array(result_list, dtype=dtype)


if __name__ == "__main__":
    pass
    # data = cal_product_total()
    # tool.get_csv(data, "product_total_quantity.csv")
